`include "define.sv"

// 数据分割器模块：将128位输入数据根据数据类型分割为32位输出
module data_splitter (
    input  wire         clk,
    input  wire         rst_n,
    
    // 上游FIFO输入接口
    input  wire [127:0] fifo_data_in,    // 128位输入数据
    output  reg         fifo_ready_out,  // 通知FIFO可以写入新数据
    input  wire         fifo_empty, //fifo空信号
    
    // 数据类型选择
    input  wire [1:0]   data_type,       // 数据类型：INT4/INT8/FP16/FP32
    
    // 下游输出接口
    output reg  [31:0]  data_out,        // 32位输出数据
    output reg          data_valid_out   // 输出数据有效信号
);

    // 内部寄存器定义
    reg [127:0] data_reg;           // 存储当前处理的128位数据
    reg [1:0]   state;              // 状态机状态
    reg [5:0]   element_count;      // 当前已输出的元素计数
    reg [5:0]   total_elements;     // 当前数据类型下总元素数
    reg [4:0]   element_width;      // 当前数据类型下单元素宽度
    
    // 状态机状态定义
    parameter RECEIVE    = 2'b00;   // 等待接收新数据
    parameter OUTPUT     = 2'b01;   // 正在输出分割后的数据
    parameter WAIT_READY = 2'b10;   // 等待下一个数据周期
    
    // 根据数据类型设置元素总数和宽度
    always @(*) begin
        case (data_type)
            `INT4: begin
                total_elements = 6'd32;  // 128/4 = 32个4位元素
                element_width = 5'd4;    // 4位宽
            end
            
            `INT8: begin
                total_elements = 6'd16;  // 128/8 = 16个8位元素
                element_width = 5'd8;    // 8位宽
            end
            
            `FP16: begin
                total_elements = 6'd8;   // 128/16 = 8个16位元素
                element_width = 5'd16;   // 16位宽
            end
            
            `FP32: begin
                total_elements = 6'd4;   // 128/32 = 4个32位元素
                element_width = 5'd5;    // 这里写5仅作占位，实际用不到
            end
            
            default: begin
                total_elements = 6'd4;
                element_width = 5'd5;
            end
        endcase
    end
    
    // 元素提取逻辑：根据数据类型和计数提取当前元素
    reg [31:0] current_element;
    always @(*) begin
        case (data_type)
            `INT4: begin
                // 提取4位元素，高位补零
                current_element = {28'b0, data_reg[127 - element_count*4 -: 4]};
            end
            
            `INT8: begin
                // 提取8位元素，高位补零
                current_element = {24'b0, data_reg[127 - element_count*8 -: 8]};
            end
            
            `FP16: begin
                // 提取16位元素，高位补零
                current_element = {16'b0, data_reg[127 - element_count*16 -: 16]};
            end
            
            `FP32: begin
                // 提取32位元素
                current_element = data_reg[127 - element_count*32 -: 32];
            end
            
            default: begin
                current_element = 32'b0;
            end
        endcase
    end
    
    // 主状态机：控制数据接收、分割和输出
    always @(posedge clk or negedge rst_n) begin
        if (!rst_n) begin
            state <= RECEIVE;
            data_reg <= 128'b0;
            data_out <= 32'b0;
            data_valid_out <= 1'b0;
            element_count <= 6'b0;
            fifo_ready_out <= 1'b1;  // 上电后允许FIFO写入
        end else begin
            case (state)
                RECEIVE: begin
                    // 只有当 fifo_empty 为 0 时，才允许接收数据
                    if (~fifo_empty) begin
                        fifo_ready_out <= 1'b1;  // 允许接收数据
                        
                        // 当有数据输入时，进入处理状态
                        if (fifo_ready_out) begin
                            data_reg <= fifo_data_in;      // 采样输入数据
                            element_count <= 6'b0;         // 元素计数清零
                            state <= OUTPUT;               // 进入输出状态
                            data_valid_out <= 1'b0;
                            fifo_ready_out <= 1'b0;        // 暂时不允许新数据写入
                        end
                    end else begin
                        fifo_ready_out <= 1'b0;  // 阻止接收数据
                    end
                end
                
                OUTPUT: begin
                    // 输出当前元素
                    data_out <= current_element;
                    
                    // 判断是否还有剩余元素
                    if (element_count + 1 < total_elements) begin
                        data_valid_out <= 1'b1;        // 输出有效
                        element_count <= element_count + 1;
                    end else begin
                        data_valid_out <= 1'b1;        // 最后一个元素也输出有效
                        element_count <= 6'b0;
                        state <= WAIT_READY;           // 进入等待状态
                    end
                end
                
                WAIT_READY: begin
                    // 输出无效，准备接收新数据
                    data_valid_out <= 1'b0;
                    fifo_ready_out <= 1'b1;            // 通知FIFO可以写入新数据
                    state <= RECEIVE;
                end
            endcase
        end
    end

endmodule 